"""
数据库迁移工具
用于处理数据库表结构变更
"""

import os
import sys
from sqlalchemy import create_engine, text
from src.config.settings import settings

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

def add_category_column():
    """为sessions表添加category字段"""
    # 使用getattr函数访问database属性以避免Pylance报错
    database_url = getattr(settings, 'DATABASE_URL', None) or getattr(settings, 'database').URL
    engine = create_engine(database_url)
    
    # 检查sessions表是否存在category列
    check_column_sql = """
    SELECT COUNT(*) as column_exists
    FROM pragma_table_info('sessions') 
    WHERE name = 'category'
    """ if 'sqlite' in database_url.lower() else """
    SELECT COUNT(*) as column_exists
    FROM information_schema.columns 
    WHERE table_name = 'sessions' AND column_name = 'category'
    """
    
    with engine.connect() as conn:
        result = conn.execute(text(check_column_sql))
        fetch_result = result.fetchone()
        column_exists = fetch_result[0] > 0 if fetch_result else False
        
        if not column_exists:
            # 添加category列
            add_column_sql = """
            ALTER TABLE sessions ADD COLUMN category VARCHAR(20) DEFAULT 'default'
            """
            try:
                conn.execute(text(add_column_sql))
                conn.commit()
                print("成功为sessions表添加category字段")
                return True
            except Exception as e:
                print(f"添加category字段失败: {e}")
                # 对于SQLite，如果表已经有数据，可能需要重建表
                if 'sqlite' in database_url.lower():
                    print("对于SQLite数据库，可能需要重建表结构")
                return False
        else:
            print("sessions表已包含category字段")
            return True

def migrate_database():
    """执行数据库迁移"""
    print("开始数据库迁移...")
    success = add_category_column()
    if success:
        print("数据库迁移完成")
    else:
        print("数据库迁移失败")
    return success

if __name__ == "__main__":
    migrate_database()